# coding:utf-8

import os, shutil
from xml.etree.ElementTree import ElementTree


def do_rename(img, start_id, i, root_path, out_path, format):
    old_img_name = img
    new_img_name = '%0{}d'.format(format) % (start_id + i) + '.jpg'
    old_img_path = os.path.join(root_path, img)
    new_img_path = os.path.join(out_path, new_img_name)

    old_xml_name = old_img_name.split('.jpg')[0] + '.xml'
    old_xml_path = os.path.join(root_path, old_xml_name)
    tree = ElementTree()
    tree.parse(old_xml_path)
    file_node = tree.find('filename')
    file_node.text = new_img_name
    path_node = tree.find('path')
    if path_node != None:
        path_text = path_node.text
        temp_list = path_text.split('\\')
        temp_list[-1] = new_img_name
        new_path_text = ''
        for i in range(len(temp_list)):
            if i == len(temp_list) - 1:
                new_path_text = new_path_text + temp_list[i]
            else:
                new_path_text = new_path_text + temp_list[i] + '\\'
        path_node.text = new_path_text
    new_xml_name = new_img_name.split('.jpg')[0] + '.xml'
    old_xml_path = os.path.join(root_path, old_xml_name)
    new_xml_path = os.path.join(out_path, new_xml_name)
    tree.write(new_xml_path, encoding="utf-8", xml_declaration=True)

    shutil.copy(old_img_path, new_img_path)
    # print('Convert {old_img} to {new_img}'.format(old_img=old_img_path, new_img=new_img_path))
    # print('Convert {old_xml} to {new_xml}'.format(old_xml=old_xml_path, new_xml=new_xml_path))


def rename_dir(start_id=0,
               root_path='./DJI_IMG0008-8',
               out_path='./out/',
               format=len('00000000')):
    '''
    start_id = 700  # 起始id
    root_path = './images'  # 图像路径
    format = len('00000000')  # 命名格式，'0000'表示'0000.jpg'
    '''
    if os.path.exists(out_path) is False:
        os.makedirs(out_path)
    lists = os.listdir(root_path)
    img_lists = [i for i in lists if i[-4:] == '.jpg']
    xml_lists = [i for i in lists if i[-4:] == '.xml']
    for i, img in enumerate(img_lists):
        try:
            do_rename(img, start_id, i, root_path, out_path, format)
        except Exception as ex:
            print("processing {} exception:{}".format(os.path.join(root_path, img), str(ex)))
    return start_id + i


def case1():
    rename_dir()


def case2(in_dir='/home/leo/Downloads/datas/road_car/all_set4k/',
          out_dir='../train_data/mix/'):
    dir_num = len(os.listdir(in_dir))
    start_id = 0
    for i, dir in enumerate(os.listdir(in_dir)):
        print(" processing dir {} of {}".format(i, dir_num))
        start_id = rename_dir(start_id=start_id, root_path=os.path.join(in_dir, dir),
                              out_path=out_dir)


if __name__ == '__main__':
    # case1()
    case2()
